import numpy as np
import os
from tqdm import tqdm

def get_file_size(file_path):
    """获取文件大小并转换为合适的单位"""
    size_bytes = os.path.getsize(file_path)
    for unit in ['B', 'KB', 'MB', 'GB']:
        if size_bytes < 1024.0:
            return f"{size_bytes:.2f} {unit}"
        size_bytes /= 1024.0
    return f"{size_bytes:.2f} TB"

def find_valid_boundary(data):
    """
    寻找数据的有效边界，使用nan值判断
    
    Args:
        data (numpy.ndarray): 输入数据
        
    Returns:
        tuple: (row_min, row_max, col_min, col_max) 有效区域的边界
        numpy.ndarray: 有效区域的掩码
    """
    # 对每个时间步的数据进行求和，得到每个位置是否有有效值
    data_sum = np.nansum(data, axis=0)
    
    # 打印数据的基本信息
    print(f"数据形状: {data.shape}")
    print(f"数据总和: {np.nansum(data)}")
    print(f"非nan值数量: {np.count_nonzero(~np.isnan(data))}")
    print(f"数据范围: [{np.nanmin(data)}, {np.nanmax(data)}]")
    
    # 创建掩码（非nan的位置为True）
    mask = ~np.isnan(data_sum)
    
    # 找到非nan值的行和列
    non_nan_rows = np.where(np.sum(mask, axis=1) > 0)[0]
    non_nan_cols = np.where(np.sum(mask, axis=0) > 0)[0]
    
    # 检查是否找到非nan值
    if len(non_nan_rows) == 0 or len(non_nan_cols) == 0:
        print("警告：未找到有效值！")
        # 返回整个数据范围作为有效区域
        return (0, data.shape[1]-1, 0, data.shape[2]-1), mask
    
    # 计算边界
    row_min, row_max = non_nan_rows[0], non_nan_rows[-1]
    col_min, col_max = non_nan_cols[0], non_nan_cols[-1]
    
    print(f"找到的有效行数: {len(non_nan_rows)}")
    print(f"找到的有效列数: {len(non_nan_cols)}")
    
    return (row_min, row_max, col_min, col_max), mask

def filter_zero_values(input_file, output_file, threshold=0.03):
    """
    过滤0值超过阈值的时间步，使用有效边界掩码计算统计值
    
    Args:
        input_file (str): 输入npy文件路径
        output_file (str): 输出npy文件路径
        threshold (float): 0值比例阈值，默认0.5（50%）
    """
    # 加载数据
    print(f"正在加载数据: {input_file}")
    data = np.load(input_file)
    print(f"原始数据形状: {data.shape}")
    print(f"原始数据大小: {get_file_size(input_file)}")
    
    # 寻找有效边界和创建掩码
    print("\n正在寻找有效边界...")
    boundary, valid_mask = find_valid_boundary(data)
    print(f"有效边界: 行范围 [{boundary[0]}, {boundary[1]}], 列范围 [{boundary[2]}, {boundary[3]}]")
    
    # 计算有效区域大小
    valid_area_size = np.sum(valid_mask)
    print(f"有效区域大小: {valid_area_size} 像素")
    
    # 记录被删除的时间索引
    removed_indices = []
    
    # 检查每个时间步
    print("\n开始检查每个时间步的统计值（使用有效边界）...")
    valid_indices = []
    for t in tqdm(range(data.shape[0])):
        time_slice = data[t]
        
        # 只计算有效区域的值
        valid_values = time_slice[valid_mask]
        
        # 计算有效区域的统计值（忽略nan值）
        if len(valid_values) > 0:
            # 只计算非nan值
            non_nan_values = valid_values[~np.isnan(valid_values)]
            if len(non_nan_values) > 0:
                zero_ratio = np.sum(non_nan_values == 0) / len(non_nan_values)
                mean_value = np.nanmean(valid_values)
                min_value = np.nanmin(valid_values)
                max_value = np.nanmax(valid_values)
            else:
                zero_ratio = 1.0
                mean_value = np.nan
                min_value = np.nan
                max_value = np.nan
        else:
            zero_ratio = 1.0
            mean_value = np.nan
            min_value = np.nan
            max_value = np.nan
        
        # 打印每个时间步的统计信息
        if zero_ratio <= threshold:
            valid_indices.append(t)
        else:
            removed_indices.append(t)
    
    # 保存有效的时间步
    filtered_data = data[valid_indices]
    
    # 保存过滤后的数据
    np.save(output_file, filtered_data)
    
    # 打印统计信息
    print("\n处理结果统计:")
    print(f"原始时间步数: {data.shape[0]}")
    print(f"删除的时间步数: {len(removed_indices)}")
    print(f"保留的时间步数: {len(valid_indices)}")
    print(f"删除的时间索引: {removed_indices}")
    print(f"过滤后数据形状: {filtered_data.shape}")
    print(f"过滤后数据大小: {get_file_size(output_file)}")
    
    # 保存删除的时间索引和有效掩码
    indices_file = output_file.replace('.npy', '_removed_indices.npy')
    mask_file = output_file.replace('.npy', '_mask.npy')
    boundary_file = output_file.replace('.npy', '_boundary.npy')
    
    np.save(indices_file, np.array(removed_indices))
    np.save(mask_file, valid_mask)
    np.save(boundary_file, np.array(boundary))
    
    print(f"已保存删除的时间索引到: {indices_file}")
    print(f"已保存有效区域掩码到: {mask_file}")
    print(f"已保存边界信息到: {boundary_file}")

def main():
    # 设置输入输出文件
    input_file = "processed_nc/srad_2016.npy"
    output_file = "processed_nc/srad_2016_filtered.npy"
    
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    
    # 执行过滤
    filter_zero_values(input_file, output_file)

if __name__ == "__main__":
    main() 